Hugging Faceのモデル学習で、ロス関数をカスタマイズする方法
こんちには。
データアナリティクス事業本部 機械学習チームの中村です。
Hugging Faceのライブラリの使い方紹介記事第５弾です。
今回は、学習時にロス関数をカスタマイズする方法をご紹介します。
ロス関数をカスタマイズする目的
ロス関数をカスタマイズする目的には、例えば以下のようなものが挙げられます。
- 不均衡なデータに対するアプローチの１つ
- 例えば、実際のタスクに取り組む歳は不均衡なデータなども取り扱うシーンが発生します。
- 不均衡なデータに対するアプローチはいくつかありますが、ロス関数に手を入れるというのも一つの手法です。
- デフォルトと異なるロス関数を使用したい
- 単純にデフォルトとは異なるロス関数に異なるものを使いたい場合もあります。
- Hugging FaceでTrainerを使用する場合は、デフォルトは
nn.CrossEntropyLossとなります。
今回はこのような場合に必要となるロス関数のカスタマイズについて説明します。
例として、不均衡なデータに対応するために、クラスごとの重みづけを実施してロスを計算する実装を行います。
実行環境
今回はGoogle Colaboratory環境で実行しました。
ハードウェアなどの情報は以下の通りです。
- GPU: Tesla T4 (GPUメモリ16GB搭載)
- CUDA: 11.2
- メモリ: 13GB
主なライブラリのバージョンは以下となります。
- transformers: 4.24.0
- datasets: 2.6.1
インストール
transformersとdatasetsをインストールします。
!pip install transformers datasets
また事前学習モデルの依存モジュールをインストールします。
!pip install fugashi !pip install ipadic !pip install sentencepiece
ベースとするコード
今回のベースとするコードは以下のとおりです。
import random
from datasets import load_dataset, DatasetDict
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from sklearn.metrics import accuracy_score, f1_score
import torch
# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")
# 実験のためデータセットを縮小したい場合はコチラを有効化
random.seed(42)
dataset = DatasetDict({
"train" : dataset['train']\
.select(random.sample(range(dataset['train'] .num_rows), k=1000)),
"validation": dataset['validation']\
.select(random.sample(range(dataset['validation'].num_rows), k=1000)),
"test" : dataset['test']\
.select(random.sample(range(dataset['test'] .num_rows), k=1000)),
})
# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
# トークナイズ処理
def tokenize(batch):
return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)
# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
.from_pretrained(model_ckpt, num_labels=num_labels)
.to(device))
# メトリクスの定義
def compute_metrics(pred):
labels = pred.label_ids
preds = pred.predictions.argmax(-1)
f1 = f1_score(labels, preds, average="weighted")
acc = accuracy_score(labels, preds)
return {"accuracy": acc, "f1": f1}
# 学習パラメータの設定
batch_size = 16
logging_steps = len(dataset_encoded["train"]) // batch_size
model_name = "sample-text-classification-bert"
training_args = TrainingArguments(
output_dir=model_name,
num_train_epochs=2,
learning_rate=2e-5,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
weight_decay=0.01,
evaluation_strategy="epoch",
disable_tqdm=False,
logging_steps=logging_steps,
push_to_hub=False,
log_level="error",
)
# Trainerの定義
trainer = Trainer(
model=model,
args=training_args,
compute_metrics=compute_metrics,
train_dataset=dataset_encoded["train"],
eval_dataset=dataset_encoded["validation"],
tokenizer=tokenizer
)
# トレーニング実行
trainer.train()
この内容についての解説は以下の記事を参照ください。
またコード中に記載の通り、実験のみを簡易に試されたい方は以下を有効にして試されてください。
逆にフルサイズのデータで試したい方は以下をコメントアウトして使用ください。
# 実験のためデータセットを縮小したい場合はコチラを有効化
random.seed(42)
dataset = DatasetDict({
"train" : dataset['train']\
.select(random.sample(range(dataset['train'] .num_rows), k=1000)),
"validation": dataset['validation']\
.select(random.sample(range(dataset['validation'].num_rows), k=1000)),
"test" : dataset['test']\
.select(random.sample(range(dataset['test'] .num_rows), k=1000)),
})
今回は有効化した前提で記事を書きます。
CustomTrainerの作成
公式の例
ロス関数をカスタムするには、Trainerクラスを独自に定義します。
以下にTrainerのサブクラスを使ってカスタマイズするAPIについて記載があります。
このページにちょうど、ロス関数をカスタムする例の記載があります。
class CustomTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
labels = inputs.get("labels")
# forward pass
outputs = model(**inputs)
logits = outputs.get("logits")
# compute custom loss (suppose one has 3 labels with different weights)
loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0]))
loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
return (loss, outputs) if return_outputs else loss
少しこの記述について解説をします。
関数の引数で与えられるinputsは辞書型で、このケースの場合以下のキーを持つようです。
dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'labels'])
モデルの出力であるoutputsも、このケースの場合以下のキーを持つようです。
odict_keys(['loss', 'logits'])
logitsは、ミニバッチの各クラスのlogitですので、(batch_size, クラス数)の形状となります。
こちらを参考にして、データセットのラベルの比率で重みづけをしてみます。
ラベルの比率を使った重みづけ
DataFrameに変換し、ラベルの比率を見てみます。
dataset.set_format(type="pandas")
label_ratio = dataset["train"][:].value_counts("label", sort=False, normalize=True)
dataset.reset_format()
print(label_ratio)
label 0 0.338 1 0.319 2 0.343 dtype: float64
偏りはそこまでありませんが、今回はこの逆数を重みづけに使ってみます。
(逆数では激しすぎる場合などは対数を使うものありかと思います)
逆数にするついでにPyTorchで使用できるようなtensor型、device設定に変換します。
class_weight = torch.tensor(1/label_ratio).clone().to(device, torch.float32) class_weight
tensor([2.9586, 3.1348, 2.9155], device='cuda:0')
こちらの重みづけをCustomTrainer内のnn.CrossEntropyLossに与えます。
from torch import nn
class CustomTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
labels = inputs.get("labels")
# forward pass
outputs = model(**inputs)
logits = outputs.get("logits")
# compute custom loss (suppose one has 3 labels with different weights)
loss_fct = nn.CrossEntropyLoss(weight=class_weight)
loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
return (loss, outputs) if return_outputs else loss
以下のコードで再度学習します。Trainerの代わりにCustomTrainerを使っています。
# モデル再定義
model = (AutoModelForSequenceClassification
.from_pretrained(model_ckpt, num_labels=num_labels)
.to(device))
# Trainerの定義
trainer = CustomTrainer(
model=model,
args=training_args,
compute_metrics=compute_metrics,
train_dataset=dataset_encoded["train"],
eval_dataset=dataset_encoded["validation"],
tokenizer=tokenizer
)
# トレーニング実行
trainer.train()
これでロス関数をカスタマイズした学習ができました。
まとめ
いかがでしたでしょうか？
今回はチューニングの際に必要になることも多い、ロス関数のカスタマイズについて記事にさせていただきました。 この実装方法で様々なロス関数のカスタマイズに汎用的に対応できると思います。
本記事がHugging Faceを使われる方の参考になれば幸いです。
